#!/usr/bin/env python3

import pathlib
import numpy as np
import monai
from typing import Optional, Union, Dict, List, Tuple, Callable
from numpy.typing import ArrayLike


from thesmuggler import smuggle
op = smuggle(f'{pathlib.Path(__file__).resolve().parents[2]}/data/operator.py')
pp = smuggle(f'{pathlib.Path(__file__).resolve().parents[2]}/data/pathparser.py')


def load_decathlon_datalist(
    data_list_file_path: str,
    is_segmentation: Optional[bool] = True,
    data_list_key: Optional[str] = "training",
    base_dir: Optional[str] = None,
) -> List[Dict]:
    """Load image/label paths of decathlon challenge from yaml file
    Returns a list of data items, each of which is a dict keyed by element names, for example:

    .. code-block::

        [
            {'image': '/workspace/data/chest_19.nii.gz',  'label': 0},
            {'image': '/workspace/data/chest_31.nii.gz',  'label': 1}
        ]

    """
    data_list_file_path = pathlib.Path(data_list_file_path)
    if not data_list_file_path.is_file():
        raise ValueError(f"Data list file {data_list_file_path} does not exist.")
    with open(data_list_file_path) as yaml_file:
        yaml_data = pp.yaml.load(yaml_file)
    if data_list_key not in yaml_data:
        raise ValueError(f'Data list {data_list_key} not specified in "{data_list_file_path}".')
    expected_data = yaml_data[data_list_key]
    # if data_list_key == "test":
    #     expected_data = [{"image": i} for i in expected_data]

    if base_dir is None:
        base_dir = data_list_file_path.parent
    return expected_data


class LoadImaged(monai.transforms.io.dictionary.LoadImaged):
    def __init__(
        self,
        keys: monai.config.KeysCollection,
        reader: Optional[Union[monai.data.image_reader.ImageReader, str]] = None,
        dtype = np.float32,
        meta_keys: Optional[monai.config.KeysCollection] = None,
        meta_key_postfix: Optional[str] = "meta_dict",
        overwriting: Optional[bool] = False,
        image_only: Optional[bool] = False,
        allow_missing_keys: Optional[bool] = False,
        *args: object,
        **kwargs: object,
    ) -> None:
        super().__init__(
            keys,
            reader,
            dtype,
            meta_keys,
            meta_key_postfix,
            overwriting,
            image_only,
            allow_missing_keys,
            *args,
            **kwargs,
        )
        self._loader = LoadImage(reader, image_only, dtype, *args, **kwargs)

    def __call__(self, data: Dict, reader: Optional[monai.data.image_reader.ImageReader] = None) -> Dict:
        reader = monai.data.image_reader.NibabelReader()
        return super().__call__(data, reader)


class LoadImage(monai.transforms.io.array.LoadImage):
    def __init__(
        self,
        reader: Optional[Union[monai.data.image_reader.ImageReader, str]] = None,
        image_only: Optional[bool] = False,
        dtype = np.float32,
        *args: object,
        **kwargs: object,
    ) -> None:
        super().__init__(reader, image_only, dtype, *args, **kwargs)
        self.func: Callable = kwargs.pop('func', None)

    def __call__(
        self, yamlmetadata: List, reader: Optional[monai.data.image_reader.ImageReader] = None
    ) -> Tuple[ArrayLike, Dict]:
        datapath = self.func(yamlmetadata) if self.func else yusongli_generate_datapath(yamlmetadata)
        img_array, meta_data = super().__call__(datapath, reader)

        meta_data['yaml_meta_data'] = yamlmetadata

        return img_array, meta_data


# COMMON = '/home/yusongli/Public/sda1/_dataset/shidaoai/img'
_timestamp = '20220407'
_epoch = '56'
_zoom = '1.1'
# _datapath = (
#     f'{op.COMMON}/_out/wangqifeng-spacial-precropped_96_224_224-dilated_k7_i2-net_myunetr_val-roi_zoom/'
#     f'{_timestamp}/{_epoch}/zoom-{_zoom}'
# )
_datapath = '/home/yusongli/_dataset/shidaoai/img/_out/'


def yusongli_generate_datapath(yamlmetadata: List) -> Tuple[str]:
    # ! <<< open debug yusongli
    # _tag = yamlmetadata[0]
    # _tag = 'train' if _tag == 'training' else _tag
    # _tag = 'val' if _tag == 'validation' else _tag
    # ! >>> clos debug

    # ! <<< open debug yusongli
    # _key = yamlmetadata[1]
    # _where = yamlmetadata[2]
    # _who = yamlmetadata[3]
    # _number = yamlmetadata[4]
    # _name = yamlmetadata[5]
    # ! >>> clos debug

    # ! <<< open debug yusongli
    # if _key == 'image':
    #     objpath = f'{_datapath}/wangqifeng-spacial/{_where}/{_who}/{_number}/{_name}'
    # elif _key == 'label':
    #     objpath = f'{_datapath}/wangqifeng-spacial-dilated_maskonly/{_where}/{_who}/{_number}/{_name}'
    # ! ===
    # _nn = yamlmetadata[6]
    # objpath = pp.filt_path(
    #     f'/home/yusongli/_dataset/shidaoai/img/_out/nn/DATASET/nnUNet_raw_data_base/nnUNet_raw_data/Task602_Z2/{_tag}_{_key}s',
    #     patterns=f'*Z2_{_nn}*'
    # )
    # ! >>> clos debug

    if yamlmetadata[0] == 'image':
        # result = pathlib.Path(f'/home/yusongli/Templates/yunet/nnUNet_raw/Dataset002_C_intensity1500_roi2.0/imagesTr/{yamlmetadata[1]}_0000.nii.gz')
        result = pathlib.Path(f'/home/yusongli/Templates/yunet/nnUNet_raw/Dataset002_C_intensity1500_roi2.0/imagesTr/{yamlmetadata[1]}_0000.nii.gz')
    elif yamlmetadata[0] == 'label':
        # result = pathlib.Path(f'/home/yusongli/Templates/yunet/nnUNet_raw/Dataset002_C_intensity1500_roi2.0/labelsTr/{yamlmetadata[1]}.nii.gz')
        result = pathlib.Path(f'/home/yusongli/Templates/yunet/nnUNet_raw/Dataset002_C_intensity1500_roi2.0/labelsTr/{yamlmetadata[1]}.nii.gz')
    return (str(result), ) if result.exists() else ('', )


def test_yusongli_generate_datapath(yamlmetadata: List) -> Tuple[str]:
    # test_file = load_decathlon_datalist(yamlmetadata, True, 'test')
    train_file = load_decathlon_datalist(yamlmetadata, True, 'training')
    # test = yamlmetadata['test']
    # LoadImaged(keys=["image", "label"])(train_file[0])
    # LoadImaged(keys=["image"])(train_file[0])
    LoadImaged(keys=["label"])(train_file[0])


if __name__ == '__main__':
    # yamlmetadata = '/home/yusongli/Documents/shidaoai_new_project/data/meta_data.yaml'
    # yamlmetadata2 = '/home/yusongli/Documents/shidaoai_new_project/data/meta_data2.yaml'
    # yamlmetadata3 = '/home/yusongli/Documents/shidaoai_new_project/data/meta_data3.yaml'
    # # with open(yamlmetadata3, 'r') as j:
    # #     yamlmetadata = pp.yaml.load(j)
    # # yusongli_generate_datapath(yamlmetadata['training'][0]['label'])

    # test_file = load_decathlon_datalist(yamlmetadata3, True, 'test')
    # train_file = load_decathlon_datalist(yamlmetadata3, True, 'training')
    # # test = yamlmetadata['test']
    # LoadImaged(keys=["image", "label"])(train_file[0])
    yamlmetadata = '/home/yusongli/Documents/shidaoai_new_project/src/data/Task607_CZ2.yaml'
    test_yusongli_generate_datapath(yamlmetadata)

